/*
 *  vdcipfw/main.c
 *
 *  The kernel module main program for Inter-VDC VM isolation.
 *
 *  Copyright (C) 2010-20xx CCMA/ITRI by Teddy Hsiung(TeddyBear)
 */

#include <linux/module.h>
#include <linux/version.h>
#include <linux/kernel.h>
#include <linux/skbuff.h>
#include <linux/init.h>
#include <linux/netfilter.h>
#include <linux/netfilter_bridge.h>
#include <linux/ip.h>
#include <linux/tcp.h>
#include <linux/udp.h>
#include <linux/in.h>
#include <net/netlink.h>
#include <net/sock.h>
#include <linux/string.h>
#include <linux/if_arp.h>
#include <linux/spinlock.h>
#include "tcp_filter.h"

#define GET_CMDID(cmd)		(cmd - TCP_FILTER_BASE_CTL)
#define GET_INFO_ARG_LEN	(sizeof(struct filter_getinfo))
#define GET_FILTERS_ARG_LEN	(sizeof(struct filter_get_rules))
#define GET_COUNTER_ARG_LEN	(sizeof(struct filter_getcounter))
static const unsigned char get_arglen[GET_CMDID(TCP_FILTER_SO_GET_MAX)+1] = {
	[GET_CMDID(TCP_FILTER_SO_GET_INFO)]	= GET_INFO_ARG_LEN,
	[GET_CMDID(TCP_FILTER_SO_GET_ENTRIES)]	= GET_FILTERS_ARG_LEN,
	[GET_CMDID(TCP_FILTER_SO_GET_COUNTER)]	= GET_FILTERS_ARG_LEN,
};

/* lock for service table */
//static DEFINE_RWLOCK(arp_nat_rule_list_lock);

/* number of rules */
static int num_rules = 0;
static atomic_t total_counter = ATOMIC_INIT(0);

static struct list_head tcp_filter_rule_list;
static spinlock_t rule_list_lock;

struct filter_rule_info { /* ADD, DEL */
    struct list_head list;
    __be32 d_ip;
    __be16 d_port;
};

static struct filter_rule_info *tcp_filter_rule_get(__be32 dip, __u16 dport)
{
    struct filter_rule_info *rule, *nrule;

    spin_lock_bh(&rule_list_lock);
    list_for_each_entry_safe(rule, nrule, &tcp_filter_rule_list, list) {
        if ((rule->d_ip == dip) && ((rule->d_port == dport) || (dport == 0))) {
            return rule;
        }
    }
    spin_unlock_bh(&rule_list_lock);

    return NULL;
}

static int do_add_rule(const void __user *user, unsigned int len)
{
    struct filter_rule_info *rule;
    struct filter_rule_info_user tmp;

    if (copy_from_user(&tmp, user, sizeof(struct filter_rule_info_user)) != 0) {
        printk("copy_from_user fail\n");
        return -EFAULT;
    }

    rule = tcp_filter_rule_get(tmp.d_ip, tmp.d_port);
    if (rule != NULL) {
	spin_unlock_bh(&rule_list_lock);
        return -EEXIST;
    }

    rule = kmalloc(sizeof(struct filter_rule_info), GFP_ATOMIC);
    if (rule == NULL) {
        printk(KERN_ERR "STATISTIC: do_add_rule: kmalloc failed.\n");
        return -ENOMEM;
    }
    memset(rule, 0, sizeof(struct filter_rule_info));
    rule->d_ip = tmp.d_ip;
    rule->d_port = tmp.d_port;
    INIT_LIST_HEAD(&rule->list);

    spin_lock_bh(&rule_list_lock);
    list_add(&rule->list, &tcp_filter_rule_list);
    num_rules++;
    spin_unlock_bh(&rule_list_lock);
    
    return 0;
}

static int do_del_rule(const void __user *user, unsigned int len)
{
    struct filter_rule_info *rule;
    struct filter_rule_info_user tmp;

    if (copy_from_user(&tmp, user, sizeof(struct filter_rule_info_user)) != 0) {
        return -EFAULT;
    }

    rule = tcp_filter_rule_get(tmp.d_ip, tmp.d_port);
    if (rule == NULL) {
        return -ENOENT;
    }

    list_del(&rule->list);
    kfree(rule);
    num_rules--;
    spin_unlock_bh(&rule_list_lock);
    return 0;
}

static int do_tcp_filter_set_ctl(struct sock *sk, int cmd, void __user *user, unsigned int len)
{
    int ret = 0;

    if (!capable(CAP_NET_ADMIN))
        return -EPERM;

    switch(cmd) {
        case TCP_FILTER_SO_SET_ADD:
            ret = do_add_rule(user, len);
            break;
        case TCP_FILTER_SO_SET_DEL:
            ret = do_del_rule(user, len);
            break;

        case TCP_FILTER_SO_SET_ZERO:
            atomic_set(&total_counter, 0);
            break;
        default:
            ret = -EINVAL;
    }
    return ret;
}

static inline int
__do_tcp_filter_get_rules(const struct filter_getinfo *get,
			  struct filter_get_rules __user *uptr)
{
    struct filter_rule_info *rule, *nrule;
    struct filter_rule_info_user entry;
    int ret = -1, count = 0;

    spin_lock_bh(&rule_list_lock);
    list_for_each_entry_safe(rule, nrule, &tcp_filter_rule_list, list) {
        if (count >= get->num_rules)
            goto out;

        entry.d_ip = rule->d_ip;
        entry.d_port = rule->d_port;
        if (copy_to_user(&uptr->ruletable[count], &entry, sizeof(entry))) {
            ret = -EFAULT;
            goto out;
        }
        count++;
    }

    ret = 0;
out:
    spin_unlock_bh(&rule_list_lock);
    return ret;
}

static int do_tcp_filter_get_ctl(struct sock *sk, int cmd, void __user *user, int *len)
{
    unsigned char arg[128];
    int ret = 0;

    if (!capable(CAP_NET_ADMIN))
        return -EPERM;

    if (copy_from_user(arg, user, get_arglen[GET_CMDID(cmd)]) != 0)
        return -EFAULT;

    switch (cmd) {
    case TCP_FILTER_SO_GET_ENTRIES:
    {
        struct filter_getinfo *get;
        int size;

        get = (struct filter_getinfo *)arg;
        size = sizeof(*get) + sizeof(struct filter_rule_info_user) * get->num_rules;
        if (*len != size) {
            printk(KERN_ERR "TCPFILTER: length: %u != %u\n", *len, size);
            ret = -EINVAL;
            goto out;
        }
        ret = __do_tcp_filter_get_rules(get, user);
    }
    break;
    case TCP_FILTER_SO_GET_INFO:
    {
        struct filter_getinfo info;
        info.num_rules = num_rules;
        if (copy_to_user(user, &info, sizeof(info)) != 0)
            ret = -EFAULT;
    }
    break;
    case TCP_FILTER_SO_GET_COUNTER:
    {
        struct filter_getcounter tmp;
        tmp.total_count = atomic_read(&total_counter);
        if (copy_to_user(user, &tmp, sizeof(tmp)) != 0)
            ret = -EFAULT;
    }
    break;
    default:
        ret = -EINVAL;
    }

out:
    return ret;
}

// Function called by BRIDGE hooks
unsigned int br_tcp_filter( unsigned int hooknum, struct sk_buff *skb,
                           const struct net_device *in,
                           const struct net_device *out,
                           int (*okfn)(struct sk_buff *))
{
    const struct iphdr *ih;
    struct tcphdr *thdr;
    struct udphdr *uhdr;
    struct filter_rule_info *rule;

    if (eth_hdr(skb)->h_proto != htons(ETH_P_IP)) {
        goto out;
    }
    ih = ip_hdr(skb);
    if (ih->protocol != IPPROTO_TCP) {
        if (ih->protocol == IPPROTO_UDP) {
            uhdr = (struct udphdr *)(((char *)ih) + ih->ihl*4);
            rule = tcp_filter_rule_get(ih->daddr, 0);
            if (rule) {
                atomic_inc(&total_counter);
	        spin_unlock_bh(&rule_list_lock);
                //kfree_skb(skb);
                //return NF_STOLEN;
                return NF_DROP;
            }
        }
        goto out;
    }
    thdr = (struct tcphdr *)(((char *)ih) + ih->ihl*4);
    
    // DROP filter rule
    rule = tcp_filter_rule_get(ih->daddr, thdr->dest);
    if (rule) {
        atomic_inc(&total_counter);
	spin_unlock_bh(&rule_list_lock);
        //kfree_skb(skb);
        //return NF_STOLEN;
        return NF_DROP;
    }

out:
    return NF_ACCEPT;
}

static struct nf_hook_ops br_tcp_filter_ops_req = {   //struct for hook function options
    .hook     = br_tcp_filter,
    .pf       = PF_BRIDGE,
    .hooknum  = NF_BR_POST_ROUTING,
    .priority = NF_BR_PRI_FIRST,
    .owner    = THIS_MODULE,
};

static struct nf_sockopt_ops br_tcp_filter_sockopts ={
    .pf             = PF_INET,
    .set_optmin     = TCP_FILTER_BASE_CTL,
    .set_optmax     = TCP_FILTER_SO_SET_MAX + 1,
    .set            = do_tcp_filter_set_ctl,
    .get_optmin     = TCP_FILTER_BASE_CTL,
    .get_optmax     = TCP_FILTER_SO_GET_MAX + 1,
    .get            = do_tcp_filter_get_ctl,
};

static int __init tcp_filter_init(void)
{
    int ret;

    spin_lock_init(&rule_list_lock);
    INIT_LIST_HEAD(&tcp_filter_rule_list);
	
    ret = nf_register_hook(&br_tcp_filter_ops_req);  // Register Hook
    if (ret < 0) {
        return ret;
    }
    ret = nf_register_sockopt(&br_tcp_filter_sockopts);
    if (ret < 0) {
        nf_unregister_hook(&br_tcp_filter_ops_req);
        return ret;
    }

    printk(KERN_INFO "TCPFILTER: Kernel Module Registered\n");
    return ret;
}

static void __exit tcp_filter_exit(void)
{
    struct filter_rule_info *rule_info, *n_rule_info;

    nf_unregister_hook(&br_tcp_filter_ops_req);  // UnRegister Hook
    nf_unregister_sockopt(&br_tcp_filter_sockopts);

    /* Destroy slb_info list */
    list_for_each_entry_safe (rule_info, n_rule_info, &tcp_filter_rule_list, list) {
        list_del (&rule_info->list);
        kfree(rule_info);
    }

    printk(KERN_INFO "TCPFILTER: Kernel Module Removed Successfully.\n" );
}

module_init(tcp_filter_init);
module_exit(tcp_filter_exit);

MODULE_DESCRIPTION("CCMA TCP packets filter Module");
MODULE_LICENSE("ITRI CCMA");
MODULE_ALIAS("CCMA TCP packets filter Module");
